# 聊天机器人及流式输出

import os
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_deepseek import ChatDeepSeek

os.environ["LANGCHAIN_TRACING_V3"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "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"
os.environ["LANGCHAIN_PROJECT"] = "langchain-community-demo"

deepseek_api_key = "sk-1dd16a258a73428d910d38c782e1c94f"

# 1、创建模型
model = ChatDeepSeek(
    api_key=deepseek_api_key,
    # deepseek-reasoner : DeepSeek-R1
    model="deepseek-chat",
    temperature=0.7,
    max_tokens=1024
)

# 2、准备提示语模板
prompt_template = ChatPromptTemplate.from_messages([
    SystemMessage(content="你是一个乐于助人的小助手，用{language}尽你所能回答所有问题。"),
    # 记住历史记录（key与下面保持一致）
    MessagesPlaceholder(variable_name='input_my_msg')
])

# 4、得到链
chain = prompt_template | model

# 保存聊天历史记录
chat_history = {}  # 所有用户的聊天都保存， key: session_id, value: 历史聊天记录对象


def get_session_history(session_id):
    """根据sessionId获取历史聊天记录对象"""
    if session_id not in chat_history:
        chat_history[session_id] = ChatMessageHistory()
    return chat_history[session_id]


do_message = RunnableWithMessageHistory(
    chain,
    get_session_history,
    # 默认的message_history_key是messages，这里修改为chat_history
    # message_history_key="chat_history",
    input_messages_key='input_my_msg'
)

config = {
    # 给当前会话定义一个sessionId
    'configurable': {'session_id': 'wugongzi'}
}

# 第一轮对话
resp = do_message.invoke(
    {
        "input_my_msg": [HumanMessage(content="你好，我叫王娜")],
        'language': '中文'
    },
    config=config
)
print(resp.content)

# 第二轮对话
resp1 = do_message.invoke(
    {
        "input_my_msg": [HumanMessage(content="请问一下我的名字叫什么？")],
        'language': '中文'
    },
    config=config
)
print(resp1.content)
# 如果之前输出是英语，这个要用中文的话，需要重新开启一个会话，设置一个新的session_id
# config = {
#     # 给当前会话定义一个sessionId
#     'configurable': {'session_id': 'wugongzi123'}
# }
# 第三轮对话
resp2 = do_message.stream(
    {
        "input_my_msg": [HumanMessage(content="请问一下当前市面上大模型有哪些？")],
        'language': '中文'
    },
    config=config
)
for resp in resp2:
    # 每一次resp都是一个token,输出不换行
    print(resp.content, end='')
